# main.py
import argparse
import sys
import os
import logging

# 添加当前目录到Python路径
# sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

# 添加项目根目录到Python路径（stock_traing目录）
project_root = os.path.dirname(os.path.abspath(__file__))  # stock_data_system目录
stock_taring_root = os.path.dirname(project_root)  # stock_traing目录

# 添加stock_traing目录到Python路径
if stock_taring_root not in sys.path:
    sys.path.insert(0, stock_taring_root)

from config.database import AppConfig
from services.data_service import StockDataService

def setup_logging(level=logging.INFO):
    """设置日志配置"""
    logging.basicConfig(
        level=level,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
        handlers=[
            logging.StreamHandler(sys.stdout),
            logging.FileHandler('stock_data.log')
        ]
    )

def print_banner():
    """打印欢迎横幅"""
    banner = """
    ====================================
        股票数据采集系统 (Baostock)
    ====================================
    """
    print(banner)

def main():
    """主程序"""
    parser = argparse.ArgumentParser(description='股票数据采集系统')
    parser.add_argument('--init', action='store_true', help='初始化数据库表结构')
    parser.add_argument('--update-basic', action='store_true', help='更新股票基本信息')
    parser.add_argument('--update-daily', nargs='?', const=1, type=int,metavar='DAYS', help='更新日线数据(可选:最近天数,默认1天)')
    parser.add_argument('--symbol', type=str, help='指定股票代码')
    parser.add_argument('--market', type=str, choices=['SH', 'SZ', 'BJ'], help='指定市场')
    parser.add_argument('--verbose', '-v', action='store_true', help='详细输出模式')

    # 在main.py的ArgumentParser中添加以下选项（如果还没有的话）
    parser.add_argument('--trading-days', type=int, default=100,help='要更新的交易天数，默认100个交易日')
    parser.add_argument('--recent-days', type=int, help='更新最近几天的数据（快捷方式）')

    parser.add_argument('--update-financial', action='store_true', help='更新财务数据')
    parser.add_argument('--update-dividend', action='store_true', help='更新分红数据')
    parser.add_argument('--years', type=int, nargs='+', help='指定年份列表，如：2020 2021 2022')
    parser.add_argument('--quarters', type=int, nargs='+', choices=[1,2,3,4], help='指定季度列表，如：1 2 3 4')

    args = parser.parse_args()

    # 设置日志级别
    log_level = logging.DEBUG if args.verbose else logging.INFO
    setup_logging(log_level)

    logger = logging.getLogger(__name__)
    print_banner()

    # 创建配置和服务
    try:
        config = AppConfig()
        service = StockDataService(config.database, config.baostock)
        logger.info("配置加载成功")
    except Exception as e:
        logger.error(f"配置加载失败: {e}")
        return 1

    # 检查数据库连接
    if not service.db_manager.check_connection():
        logger.error("数据库连接失败，请检查配置和PostgreSQL服务是否运行")
        return 1

    logger.info("数据库连接成功")

    try:
        # 初始化数据库
        if args.init:
            logger.info("正在初始化数据库表结构...")
            service.initialize_database()
            logger.info("数据库初始化完成")
            return 0

        # 更新股票基本信息
        if args.update_basic:
            logger.info("正在更新股票基本信息...")
            count = service.update_stock_basic()
            logger.info(f"成功更新 {count} 只股票基本信息")
            return 0

        # 更新日线数据
        if args.update_daily is not None:
            days = args.update_daily
            if args.symbol:
                symbols = [args.symbol]
                logger.info(f"正在更新 {args.symbol} 的最近 {days} 天日线数据...")
                count = service.update_daily_data(symbols, days=days)
                logger.info(f"成功更新 {args.symbol} 的 {count} 条日线数据")
            else:
                logger.info(f"正在更新所有股票的最近 {days} 天数据...")
                count = service.update_all_stocks_daily(days)
                logger.info(f"成功更新 {count} 条日线数据")
            return 0


        # 在main函数中添加对应的处理逻辑
        if args.update_financial:
            years = args.years if args.years else None
            quarters = args.quarters if args.quarters else None

            logger.info("正在更新财务数据...")
            if args.symbol:
                count = service.update_financial_data([args.symbol], years, quarters)
            else:
                count = service.update_all_financial_data(years, quarters)
            logger.info(f"成功更新 {count} 条财务数据")
            return 0

        if args.update_dividend:
            years = args.years if args.years else None

            logger.info("正在更新分红数据...")
            if args.symbol:
                count = service.update_dividend_data([args.symbol], years)
            else:
                count = service.update_all_dividend_data(years)
            logger.info(f"成功更新 {count} 条分红数据")
            return 0


        # 在main函数中修改对应的处理逻辑
        if args.update_daily is not None:
            if args.symbol:
                symbols = [args.symbol]
                logger.info(f"正在更新 {args.symbol} 的 {args.trading_days} 个交易日数据...")
                count = service.update_daily_data(symbols, trading_days=args.trading_days)
                logger.info(f"成功更新 {args.symbol} 的 {count} 条日线数据")
            else:
                if args.recent_days:
                    # 使用最近几天模式
                    logger.info(f"正在更新所有股票的最近 {args.recent_days} 天数据...")
                    count = service.update_recent_daily_data(args.recent_days)
                else:
                    # 使用交易日模式
                    logger.info(f"正在更新所有股票的 {args.trading_days} 个交易日数据...")
                    count = service.update_all_stocks_daily(args.trading_days)
                logger.info(f"成功更新 {count} 条日线数据")


        # 如果没有参数，显示帮助信息
        if len(sys.argv) == 1:
            parser.print_help()
            print("\n示例用法:")
            print("  python main.py --init                    # 初始化数据库")
            print("  python main.py --update-basic           # 更新股票列表")
            print("  python main.py --update-daily           # 更新最近1天数据")
            print("  python main.py --update-daily 5         # 更新最近5天数据")
            print("  python main.py --update-daily --symbol 000001  # 更新指定股票")
        return 0

    except KeyboardInterrupt:
        logger.info("用户中断操作")
        return 130
    except Exception as e:
        logger.error(f"执行过程中发生错误: {e}")
        return 1

if __name__ == "__main__":
    exit_code = main()
    sys.exit(exit_code)